import os
import tensorflow as tf
from keras import layers


class MyDense(layers.Layer):
    """
    自定义全连接层
    """

    def __init__(self, input_dim, output_dim):
        super(MyDense, self).__init__()
        self.kernel = self.add_weight(name="w", shape=[input_dim, output_dim], trainable=True)
        # self.kernel = self.add_weight(name="w", shape=[input_dim, output_dim], trainable=False)
        # self.kernel = tf.Variable(tf.random.normal(shape=[input_dim, output_dim]), trainable=True)
        # self.kernel = tf.Variable(tf.random.normal(shape=[input_dim, output_dim]), trainable=False)
        self.bias = self.add_weight(name="b", shape=[output_dim])

    def call(self, inputs, *args, **kwargs):
        """
        执行当前层的前向运算逻辑
        :param inputs: 输入数据
        :param args: 其他普通参数
        :param kwargs: 其它关键词参数
        :return:
        """
        # 执行前向运算逻辑，矩阵乘法和矩阵加法运算
        output = inputs @ self.kernel + self.bias
        # 执行激活函数运算
        output = tf.nn.relu(output)
        return output


if __name__ == '__main__':
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

    layer = MyDense(4, 3)
    print(layer.variables)
    print(layer.trainable_variables)
    print(layer.non_trainable_variables)

    inputs = tf.constant(tf.random.normal([1,4]))
    result = layer.call(inputs=inputs)
    print(result)
